//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift Atomics open source project
//
// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

import _AtomicsShims

// Double-wide atomic primitives on x86_64 CPUs aren't available by default
// on Linux distributions, and we cannot currently enable them automatically.
#if !(os(Linux) && arch(x86_64)) || ENABLE_DOUBLEWIDE_ATOMICS

/// A class type that supports atomic strong references.
public protocol AtomicReference: AnyObject, AtomicOptionalWrappable
where
  AtomicRepresentation == AtomicReferenceStorage<Self>,
  AtomicOptionalRepresentation == AtomicOptionalReferenceStorage<Self>
{
  // These were added as a workaround for https://bugs.swift.org/browse/SR-10251
  // FIXME: We should remove these once the package requires a
  // compiler version that contains that fix.
  override associatedtype AtomicRepresentation = AtomicReferenceStorage<Self>
  override associatedtype AtomicOptionalRepresentation =
    AtomicOptionalReferenceStorage<Self>
}

/// The maximum number of other threads that can start accessing a
/// strong reference before an in-flight update needs to cancel and
/// retry.
@inlinable @inline(__always)
internal var _concurrencyWindow: Int { 20 }

extension Unmanaged {
  internal func retain(by delta: Int) {
    _sa_retain_n(toOpaque(), UInt32(delta))
  }

  internal func release(by delta: Int) {
    _sa_release_n(toOpaque(), UInt32(delta))
  }
}

extension Unmanaged {
  fileprivate static func passRetained(_ instance: __owned Instance?) -> Self? {
    guard let instance = instance else { return nil }
    // Note: Swift 5.2 doesn't like this optional promotion without the explicit cast
    return .passRetained(instance) as Self
  }
}

extension DoubleWord {
  fileprivate init(_raw: UnsafeRawPointer?, readers: Int, version: Int) {
    let r = UInt(bitPattern: readers) & Self._readersMask
    assert(r == readers)
    self.init(
      high: r | (UInt(bitPattern: version) &<< Self._readersBitWidth),
      low: UInt(bitPattern: _raw))
  }

  @inline(__always)
  fileprivate var _raw: UnsafeMutableRawPointer? {
    get { UnsafeMutableRawPointer(bitPattern: low) }
    set { low = UInt(bitPattern: newValue) }
  }

  @inline(__always)
  fileprivate var _unmanaged: Unmanaged<AnyObject>? {
    guard let raw = _raw else { return nil }
    return .fromOpaque(raw)
  }

  @inline(__always)
  fileprivate static var _readersBitWidth: Int {
    // This reserves 8 bits for the accesses-in-flight counter on 32-bit
    // systems, and 16 bits on 64-bit systems.
    Int.bitWidth / 4
  }

  @inline(__always)
  fileprivate static var _readersMask: UInt { (1 &<< _readersBitWidth) - 1  }

  @inline(__always)
  fileprivate var _readers: Int {
    get { Int(bitPattern: high & Self._readersMask) }
    set {
      let n = UInt(bitPattern: newValue) & Self._readersMask
      assert(n == newValue)
      high = (high & ~Self._readersMask) | n
    }
  }

  @inline(__always)
  fileprivate var _version: Int {
    get { Int(bitPattern: high &>> Self._readersBitWidth) }
    set {
      // Silently truncate any high bits we cannot store.
      high = (
        (UInt(bitPattern: newValue) &<< Self._readersBitWidth) |
        high & Self._readersMask)
    }
  }
}

extension Optional where Wrapped == AnyObject {
  fileprivate var _raw: UnsafeMutableRawPointer? {
    guard let value = self else { return nil }
    return Unmanaged.passUnretained(value).toOpaque()
  }
}

extension UnsafeMutablePointer where Pointee == _AtomicReferenceStorage {
  @inlinable @inline(__always)
  internal var _extract: UnsafeMutablePointer<DoubleWord.AtomicRepresentation> {
    UnsafeMutableRawPointer(self)
      .assumingMemoryBound(to: DoubleWord.AtomicRepresentation.self)
  }
}

@usableFromInline
internal struct _AtomicReferenceStorage {
  internal typealias Storage = DoubleWord.AtomicRepresentation
  internal var _storage: Storage

  @usableFromInline
  internal init(_ value: __owned AnyObject?) {
    let dword = DoubleWord(
      _raw: Unmanaged.passRetained(value)?.toOpaque(),
      readers: 0,
      version: 0)
    _storage = Storage(dword)
  }

  @usableFromInline
  internal func dispose() -> AnyObject? {
    let value = _storage.dispose()
    precondition(value._readers == 0,
      "Attempt to dispose of a busy atomic strong reference \(value)")
    return value._unmanaged?.takeRetainedValue()
  }

  private static func _startLoading(
    from pointer: UnsafeMutablePointer<Self>,
    hint: DoubleWord? = nil
  ) -> DoubleWord {
    var old = hint ?? Storage.atomicLoad(at: pointer._extract, ordering: .relaxed)
    if old._raw == nil {
      atomicMemoryFence(ordering: .acquiring)
      return old
    }
    // Increment reader count
    while true {
      let new = DoubleWord(
        _raw: old._raw,
        readers: old._readers &+ 1,
        version: old._version)
      var done: Bool
      (done, old) = Storage.atomicWeakCompareExchange(
        expected: old,
        desired: new,
        at: pointer._extract,
        successOrdering: .acquiring,
        failureOrdering: .acquiring)
      if done { return new }
      if old._raw == nil { return old }
    }
  }

  private static func _finishLoading(
    _ value: DoubleWord,
    from pointer: UnsafeMutablePointer<Self>
  ) -> AnyObject? {
    if value._raw == nil { return nil }

    // Retain result before we exit the access.
    let result = value._unmanaged!.takeUnretainedValue()

    // Decrement reader count, using the version number to prevent any
    // ABA issues.
    var current = value
    var done = false
    repeat {
      assert(current._readers >= 1)
      (done, current) = Storage.atomicWeakCompareExchange(
        expected: current,
        desired: DoubleWord(
          _raw: current._raw,
          readers: current._readers &- 1,
          version: current._version),
        at: pointer._extract,
        successOrdering: .acquiringAndReleasing,
        failureOrdering: .acquiring)
    } while !done && current._raw == value._raw && current._version == value._version
    if !done {
      // The reference changed while we were loading it. Cancel out
      // our part of the refcount bias for the loaded instance.
      value._unmanaged!.release()
    }
    return result
  }

  @usableFromInline
  internal static func atomicLoad(
    at pointer: UnsafeMutablePointer<Self>
  ) -> AnyObject? {
    let new = _startLoading(from: pointer)
    return _finishLoading(new, from: pointer)
  }

  /// Try updating the current value from `old` to `new`. Don't do anything if the
  /// current value differs from `old`.
  ///
  /// Returns a tuple `(exchanged, original)` where `exchanged` indicates whether the
  /// update was successful. If `exchange` true, then `original` contains the original
  /// reference value; otherwise `original` is nil.
  ///
  /// On an unsuccessful exchange, this function updates `old` to the latest known value,
  /// and reenters the current thread as a reader of it.
  private static func _tryExchange(
    old: inout DoubleWord,
    new: Unmanaged<AnyObject>?,
    at pointer: UnsafeMutablePointer<Self>
  ) -> (exchanged: Bool, original: AnyObject?) {
    let new = DoubleWord(_raw: new?.toOpaque(), readers: 0, version: old._version &+ 1)
    guard let ref = old._unmanaged else {
      // Try replacing the current nil value with the desired new value.
      let (done, current) = Storage.atomicCompareExchange(
        expected: old,
        desired: new,
        at: pointer._extract,
        ordering: .acquiringAndReleasing)
      if done {
        return (true, nil)
      }
      // Someone else changed the reference. Give up for now.
      old = _startLoading(from: pointer, hint: current)
      return (false, nil)
    }
    assert(old._readers >= 1)
    var delta = old._readers + _concurrencyWindow
    ref.retain(by: delta)
    while true {
      // Try replacing the current value with the desired new value.
      let (done, current) = Storage.atomicCompareExchange(
        expected: old,
        desired: new,
        at: pointer._extract,
        ordering: .acquiringAndReleasing)
      if done {
        // Successfully replaced the object. Clean up extra retains.
        assert(current._readers == old._readers)
        assert(current._raw == old._raw)
        assert(current._readers <= delta)
        ref.release(by: delta - current._readers + 1) // +1 is for our own role as a reader.
        return (true, ref.takeRetainedValue())
      }
      if current._version != old._version {
        // Someone else changed the reference. Give up for now.
        ref.release(by: delta + 1) // +1 covers our reader bias
        old = _startLoading(from: pointer, hint: current)
        return (false, nil)
      }
      // Some readers entered or left while we were processing things. Try again.
      assert(current._raw == old._raw)
      assert(current._raw != nil)
      assert(current._readers >= 1)
      if current._readers > delta {
        // We need to do more retains to cover readers.
        let d = current._readers + _concurrencyWindow
        ref.retain(by: d - delta)
        delta = d
      }
      old = current
    }
  }

  @usableFromInline
  internal static func atomicExchange(
    _ desired: __owned AnyObject?,
    at pointer: UnsafeMutablePointer<Self>
  ) -> AnyObject? {
    let new = Unmanaged.passRetained(desired)
    var old = _startLoading(from: pointer)
    var (exchanged, original) = _tryExchange(old: &old, new: new, at: pointer)
    while !exchanged {
      (exchanged, original) = _tryExchange(old: &old, new: new, at: pointer)
    }
    return original
  }

  @usableFromInline
  internal static func atomicCompareExchange(
    expected: AnyObject?,
    desired: __owned AnyObject?,
    at pointer: UnsafeMutablePointer<Self>
  ) -> (exchanged: Bool, original: AnyObject?) {
    let expectedRaw = expected._raw
    let new = Unmanaged.passRetained(desired)
    var old = _startLoading(from: pointer)
    while old._raw == expectedRaw {
      let (exchanged, original) = _tryExchange(old: &old, new: new, at: pointer)
      if exchanged {
        assert(original === expected)
        return (true, original)
      }
    }
    // We did not find the expected value. Cancel the retain of the new value
    // and return the old value.
    new?.release()
    return (false, _finishLoading(old, from: pointer))
  }
}

@frozen
public struct AtomicReferenceStorage<Value: AnyObject> {
  @usableFromInline
  internal var _storage: _AtomicReferenceStorage

  @inlinable
  public init(_ value: __owned Value) {
    _storage = .init(value)
  }

  @inlinable
  public func dispose() -> Value {
    return unsafeDowncast(_storage.dispose()!, to: Value.self)
  }
}

extension AtomicReferenceStorage {
  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  static func _extract(
    _ ptr: UnsafeMutablePointer<Self>
  ) -> UnsafeMutablePointer<_AtomicReferenceStorage> {
    // `Self` is layout-compatible with its only stored property.
    return UnsafeMutableRawPointer(ptr)
      .assumingMemoryBound(to: _AtomicReferenceStorage.self)
  }
}

extension AtomicReferenceStorage: AtomicStorage {
  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicLoad(
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicLoadOrdering
  ) -> Value {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicLoad(at: Self._extract(pointer))
    return unsafeDowncast(result!, to: Value.self)
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicStore(
    _ desired: __owned Value,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicStoreOrdering
  ) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    _ = _AtomicReferenceStorage.atomicExchange(
      desired,
      at: _extract(pointer))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicExchange(
    _ desired: __owned Value,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicUpdateOrdering
  ) -> Value {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicExchange(
      desired,
      at: _extract(pointer))
    return unsafeDowncast(result!, to: Value.self)
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicCompareExchange(
    expected: Value,
    desired: __owned Value,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicUpdateOrdering
  ) -> (exchanged: Bool, original: Value) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    return (result.exchanged, unsafeDowncast(result.original!, to: Value.self))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicCompareExchange(
    expected: Value,
    desired: __owned Value,
    at pointer: UnsafeMutablePointer<Self>,
    successOrdering: AtomicUpdateOrdering,
    failureOrdering: AtomicLoadOrdering
  ) -> (exchanged: Bool, original: Value) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    return (result.exchanged, unsafeDowncast(result.original!, to: Value.self))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicWeakCompareExchange(
    expected: Value,
    desired: __owned Value,
    at pointer: UnsafeMutablePointer<Self>,
    successOrdering: AtomicUpdateOrdering,
    failureOrdering: AtomicLoadOrdering
  ) -> (exchanged: Bool, original: Value) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    return (result.exchanged, unsafeDowncast(result.original!, to: Value.self))
  }
}

@frozen
public struct AtomicOptionalReferenceStorage<Instance: AnyObject> {
  @usableFromInline
  internal var _storage: _AtomicReferenceStorage

  @inlinable
  public init(_ value: __owned Instance?) {
    _storage = .init(value)
  }

  @inlinable
  public func dispose() -> Instance? {
    guard let value = _storage.dispose() else { return nil }
    return unsafeDowncast(value, to: Instance.self)
  }
}

extension AtomicOptionalReferenceStorage {
  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  static func _extract(
    _ ptr: UnsafeMutablePointer<Self>
  ) -> UnsafeMutablePointer<_AtomicReferenceStorage> {
    // `Self` is layout-compatible with its only stored property.
    return UnsafeMutableRawPointer(ptr)
      .assumingMemoryBound(to: _AtomicReferenceStorage.self)
  }
}

extension AtomicOptionalReferenceStorage: AtomicStorage {
  public typealias Value = Instance?

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicLoad(
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicLoadOrdering
  ) -> Instance? {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicLoad(at: Self._extract(pointer))
    guard let r = result else { return nil }
    return unsafeDowncast(r, to: Instance.self)
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicStore(
    _ desired: __owned Instance?,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicStoreOrdering
  ) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    _ = _AtomicReferenceStorage.atomicExchange(
      desired,
      at: _extract(pointer))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicExchange(
    _ desired: __owned Instance?,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicUpdateOrdering
  ) -> Instance? {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicExchange(
      desired,
      at: _extract(pointer))
    guard let r = result else { return nil }
    return unsafeDowncast(r, to: Instance.self)
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicCompareExchange(
    expected: Instance?,
    desired: __owned Instance?,
    at pointer: UnsafeMutablePointer<Self>,
    ordering: AtomicUpdateOrdering
  ) -> (exchanged: Bool, original: Instance?) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    guard let original = result.original else { return (result.exchanged, nil) }
    return (result.exchanged, unsafeDowncast(original, to: Instance.self))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicCompareExchange(
    expected: Instance?,
    desired: __owned Instance?,
    at pointer: UnsafeMutablePointer<Self>,
    successOrdering: AtomicUpdateOrdering,
    failureOrdering: AtomicLoadOrdering
  ) -> (exchanged: Bool, original: Instance?) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    guard let original = result.original else { return (result.exchanged, nil) }
    return (result.exchanged, unsafeDowncast(original, to: Instance.self))
  }

  @inlinable @inline(__always)
  @_alwaysEmitIntoClient
  @_semantics("atomics.requires_constant_orderings")
  public static func atomicWeakCompareExchange(
    expected: Instance?,
    desired: __owned Instance?,
    at pointer: UnsafeMutablePointer<Self>,
    successOrdering: AtomicUpdateOrdering,
    failureOrdering: AtomicLoadOrdering
  ) -> (exchanged: Bool, original: Instance?) {
    // FIXME: All orderings are treated as acquiring-and-releasing.
    let result = _AtomicReferenceStorage.atomicCompareExchange(
      expected: expected,
      desired: desired,
      at: _extract(pointer))
    guard let original = result.original else { return (result.exchanged, nil) }
    return (result.exchanged, unsafeDowncast(original, to: Instance.self))
  }
}

#endif // ENABLE_DOUBLEWIDE_ATOMICS
